Batch Rebalance Data Sampler#47340
Conversation
New sampler that balances sequence length cost across DP ranks while
minimizing padding waste. Uses Poorman's Batch Rebalance algorithm with
configurable cost model.
Changes:
- trainer_pt_utils.py: Add BatchRebalanceSampler class
- training_args.py: Add 'batch_rebalance' to train_sampling_strategy
choices, batch_rebalance_alpha and batch_rebalance_max_tokens params
- trainer.py: Add batch_rebalance branch in _get_train_sampler
Usage:
TrainingArguments(train_sampling_strategy='batch_rebalance',
batch_rebalance_alpha=0.001)
Tests cover: - Basic correctness (coverage, micro-batch count) - No overlap between ranks within same global batch - Effective batch size invariant - max_tokens constraint enforcement - Determinism (same seed = same result) - Custom cost_fn override
PyTorch's Sampler base class doesn't accept data_source kwarg. Drop the super().__init__() call entirely (consistent with other HF samplers like LengthGroupedSampler).
…ens rebalance - Remove micro_batch_size from BatchRebalanceSampler: it was stored but never read anywhere in the class, and its name was misleading since it did not reflect the actual minimum per-group size used by the algorithm. - Fix a bug in _poormans_rebalance where counts[gi] -= 1 was applied unconditionally before searching for a donor group, causing samples to be silently dropped when max_tokens could not be satisfied and no group had spare capacity. Now the decrement only happens once a donor is found, and a warning is logged (with a fallback to exceed max_tokens rather than drop samples) when no donor is available. - Fix docstring typo (Poorman's -> Poor Man's). - Update/add tests: drop micro_batch_size from existing tests, add a regression test for the max_tokens data-loss fix, and rewrite the cost_fn test to directly assert on _cost() output instead of an unstable partition overlap comparison.
…wiring BatchRebalanceSampler yields a full batch of sample indices per iteration (BatchSampler semantics), but _get_dataloader() was passing it via DataLoader's sampler= kwarg together with batch_size=, which expects the sampler to yield one sample index at a time. This caused 'TypeError: list indices must be integers or slices, not list' as soon as train_sampling_strategy="batch_rebalance" was used with the Trainer. Also fixes a NameError (BatchEncoding was referenced in _get_train_sampler but never imported), which was the first failure hit before reaching the DataLoader wiring bug. Together these two bugs meant the batch_rebalance strategy could never actually run end-to-end through Trainer.train(); verified with a CPU multi-process (torchrun --nproc_per_node=4) benchmark on Qwen2.5-0.5B with a synthetic variable-length dataset after the fix.
When n < K * min_per_group, there are not enough samples to fill all groups. Previously this was silently patched by trimming counts (and a dead 'total < n' branch that could never trigger). Now it raises a clear error telling the user to reduce effective_batch_size or increase dataset size.
- Replace can_increase(cts, gi) with can_transfer(cts, from_idx, to_idx) which simulates the actual -1/+1 transfer to compute the correct start index and max_len for the token limit check. The previous implementation used stale positions that didn't account for shifts caused by the donor group shrinking. - When the lowest-cost group cannot accept a transfer due to max_tokens, iterate through remaining candidates by ascending cost instead of falling through to the nail-house branch or breaking entirely.
CI recapDashboard: View test results in Grafana |
Each iteration of the rebalance loop moves exactly one sample between groups, so worst-case convergence is O(n). The fixed `K * 30` cap could therefore truncate before convergence for large `effective_batch_size` with a small `dp_size * grad_accum` (K), leaving the partition at a non-optimal cost spread. Use `max(K * 30, effective_batch_size)` as the budget: the `effective_batch_size` term covers the move-dominated regime (large batch, small K), while `K * 30` remains a floor for the detection-dominated regime (tiny batch, n ~= K) where convergence is reached via `spread == 0` / `seen` rather than the move budget. `best_counts` is still tracked, so hitting the cap only degrades the result gracefully (no samples are dropped). Adds a regression test (large n=246, small K=2, heavily skewed lengths) that asserts full sample coverage across all ranks and convergence to the brute-force-optimal contiguous K=2 cost spread.
GPU DDP benchmarkRan the sampler end-to-end through the Setup
Results
Step 5 sample distribution (step 5 = Reproduce: benchmark script is at https://gist.github.com/delock/516efc09d9a4217e209729e91b9294de ( |
|
| per_device_bs | group_by_length | batch_rebalance |
|---|---|---|
| 4 | 17.7 GB | 12.3 GB |
| 6 | 24.5 GB | 15.4 GB |
| 8 | 31.3 GB | 20.4 GB |
| 10 | OOM | 21.7 GB |
| 12 | OOM | 24.5 GB |
| 14 | OOM | 27.1 GB |
| 16 | OOM | ~30 GB (borderline) |
| 18 | OOM | OOM |
group_by_lengthOOMs atper_device_bs ≥ 10→ largest reliable batch = 8 (peak ~31 GB, right at the edge).batch_rebalancefits reliably up to 14 (peak ~28–30 GB); 16 is borderline (fits on short runs, OOMs over a full epoch); ≥18 OOMs.
→ batch_rebalance tolerates roughly 2× the per-device batch before OOM (8 → 14 reliable, 16 borderline; effective batch 8×2×2 = 32 vs 14×2×2 = 56), and at every batch size its peak is ~30–40% lower.
Why: group_by_length globally sorts by length, so the longest sequences land in a single micro-batch (per_device × 1024), and that one batch sets peak VRAM (the LM-head logits + activations dominate). batch_rebalance balances cost within each global batch, so the long-sequence group shrinks to just 2–3 samples → padded-tokens-per-micro-batch is bounded → no single batch spikes.
2) Throughput when each runs at its own max batch
| sampler | max per_device_bs | real tok/s | peak VRAM |
|---|---|---|---|
| group_by_length | 8 | 30,174 | 32.3 GB |
| batch_rebalance | 14 | 37,957 | 29.6 GB |
→ batch_rebalance delivers ~1.26× higher real-tokens/s (37,957 vs 30,174) — and does so at lower peak VRAM (29.6 vs 32.3 GB).
The gain is not from per-step efficiency (at equal per_device_bs the two are close). It comes from batch_rebalance being able to use a larger batch on the same GPU → better GPU occupancy → more real tokens/s. group_by_length is memory-capped at per_device_bs=8 by its long-batch spike; batch_rebalance converts the extra memory headroom into throughput.
Net: compared to group_by_length, batch_rebalance has lower peak VRAM at every batch size, tolerates ~2× the batch before OOM, and when each is run at its max batch it's ~1.26× faster in real tokens/s.
|
Thank you for your contribution 🤗! CI Security Gate — automatic approval blockedThis PR was not automatically approved for CI because the security gate failed. Possible reasons:
See the workflow run for the exact violations. A maintainer can review and manually approve CI if a finding is a false positive. |
1 participant
What does this PR do?
This PR introduce a training data sampler using a batch rebalance technique that reduces padding, balance cost between ranks and grad accumulation steps, which improves training efficiency.
Different from packed sequence data sampling, batch rebalance data sampling does not require varlen support in attention implementation. This makes it idea for researching or hardware that does not have varlen attention kernel implementation yet.
Different from group length sampler, this sampler does not require a global sort, so the sampler maintains true randomness in each effective batch.
How it works
The basic idea of batch rebalance data sampler is like this:
In distributed training, each rank processes a subset of samples called a micro-batch. When samples have different sequence lengths, the rank that happens to get longer sequences takes more time to compute, while other ranks finish early and wait. This idle waiting wastes GPU time.
Batch rebalance solves this by looking at all samples in a global batch, sorting them by length, and distributing them across ranks so that each rank gets a roughly equal computational cost — not an equal number of samples, but an equal workload. Short-sequence ranks get more samples; long-sequence ranks get fewer. The assignment is recomputed for every global batch using a lightweight greedy heuristic, with no communication between ranks.
For example, say you have 2 ranks, gradient accumulation of 1, and a global batch of 8 samples with token lengths [128, 256, 512, 768, 1024, 1536, 2048, 4096]. A random split assigns 4 samples per rank. If rank 0 happens to get [128, 512, 1024, 4096], every sample in that micro-batch is padded to 4096 — even the 128-token one — so rank 0 computes 4×4096 = 16,384 padded tokens. Meanwhile rank 1 gets [256, 768, 1536, 2048], padding to only 2048, computing 4×2048 = 8,192 tokens. Rank 0 takes twice as long, and rank 1 sits idle.
Batch rebalance sorts all 8 samples and assigns them in unequal counts: rank 0 gets just the 2 longest samples [4096, 2048] (2×4096 = 8,192 tokens), while rank 1 gets the remaining 6 shorter samples [1536, 1024, 768, 512, 256, 128] (6×1536 = 9,216 tokens). The workload is now nearly equal — no idle waiting.
Batch rebalance sampler uses a cost model to estimate the compute of each group of samples, then balances the total cost across ranks within the same synchronization step.
Specifically, for each global batch, the sampler sorts samples by length, partitions them into dp_size × grad_accum groups, and assigns groups to ranks so that ranks sharing the same gradient accumulation slot receive similar costs. Longer-sequence groups contain fewer samples; shorter-sequence groups contain more — but the estimated compute is even.
The default cost model is intercept + bs × max_len + alpha × bs × max_len², where bs is the number of samples in the group and max_len is the longest sequence. The quadratic term captures the super-linear scaling of attention. Users can plug in their own cost function via the cost_fn parameter to match their model's actual compute characteristics.
How to use
Set
--train_sampling_strategy batch_rebalancein your training command. No other changes to your dataset or model code are needed.python train.py
--train_sampling_strategy batch_rebalance
--per_device_train_batch_size 4
--gradient_accumulation_steps 2
Optional tuning parameters:
I confirm that this PR description and code is written with aid of LLM, and reviewed by a human being prior to PR submission.
Before submitting
[ ] (First-time contributors only): I confirm that this PR description and code is not written by an LLM or code agent
Pull Request checks?
to it if that's the case.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.